package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

// orDone 函数接收一个done channel和一个输入channel
// 它返回一个新channel，当done收到信号或输入channel收到信号时，返回的channel会收到通知
func orDone(done <-chan struct{}, c <-chan interface{}) <-chan interface{} {
	valStream := make(chan interface{})

	go func() {
		defer close(valStream)

		for {
			select {
			case <-done:
				// 如果收到取消信号，立即返回
				return
			case v, ok := <-c:
				if !ok {
					// 如果输入channel关闭，也立即返回
					return
				}

				// 尝试发送值，但也接收取消信号
				select {
				case valStream <- v:
				case <-done:
					return
				}
			}
		}
	}()

	return valStream
}

// 创建一个可以被取消的数据流生成器
func generator(done <-chan struct{}, integers ...int) <-chan interface{} {
	intStream := make(chan interface{})

	go func() {
		defer close(intStream)

		for _, i := range integers {
			// 在每次发送前检查是否被取消
			select {
			case <-done:
				return
			case intStream <- i:
			}

			// 模拟一些处理时间
			time.Sleep(1 * time.Second)
		}
	}()

	return intStream
}

// primeFinder 找出质数
func primeFinder(done <-chan struct{}, intStream <-chan interface{}) <-chan interface{} {
	primeStream := make(chan interface{})

	go func() {
		defer close(primeStream)

		// 使用orDone来处理取消和输入关闭的情况
		for val := range orDone(done, intStream) {
			integer, ok := val.(int)
			if !ok {
				continue
			}

			// 检查是否为质数
			if isPrime(integer) {
				select {
				case <-done:
					return
				case primeStream <- integer:
				}
			}
		}
	}()

	return primeStream
}

// 简单的质数检查函数
func isPrime(n int) bool {
	if n <= 1 {
		return false
	}
	if n <= 3 {
		return true
	}
	if n%2 == 0 || n%3 == 0 {
		return false
	}

	i := 5
	for i*i <= n {
		if n%i == 0 || n%(i+2) == 0 {
			return false
		}
		i += 6
	}
	return true
}

// 演示使用fanIn模式合并多个channel，并支持取消
func fanIn(done <-chan struct{}, channels ...<-chan interface{}) <-chan interface{} {
	var wg sync.WaitGroup
	multiplexedStream := make(chan interface{})

	// 为每个channel启动一个goroutine
	multiplex := func(c <-chan interface{}) {
		defer wg.Done()

		// 使用orDone处理取消
		for i := range orDone(done, c) {
			select {
			case <-done:
				return
			case multiplexedStream <- i:
			}
		}
	}

	// 收集所有输入channel
	wg.Add(len(channels))
	for _, c := range channels {
		go multiplex(c)
	}

	// 等待所有输入处理完成后关闭输出
	go func() {
		wg.Wait()
		close(multiplexedStream)
	}()

	return multiplexedStream
}

// 可取消的worker，会在指定时间后完成
func worker(done <-chan struct{}, id int, delay time.Duration) <-chan interface{} {
	c := make(chan interface{}, 1)

	go func() {
		defer close(c)

		fmt.Printf("Worker %d: 开始工作，预计需要 %v\n", id, delay)

		select {
		case <-done:
			fmt.Printf("Worker %d: 被取消\n", id)
			return
		case <-time.After(delay):
			// 工作完成
		}

		// 工作完成后发送结果
		select {
		case <-done:
			fmt.Printf("Worker %d: 完成，但结果被丢弃（已取消）\n", id)
		case c <- id:
			fmt.Printf("Worker %d: 工作完成!\n", id)
		}
	}()

	return c
}

// 演示超时机制
func withTimeout(timeout time.Duration) (<-chan struct{}, func()) {
	done := make(chan struct{})

	// 启动一个超时计时器
	timer := time.AfterFunc(timeout, func() {
		close(done)
	})

	// 返回done channel和一个取消函数
	// 取消函数可以在超时前停止操作
	cancel := func() {
		timer.Stop()
		close(done)
	}

	return done, cancel
}

func main() {
	// 设置随机数种子
	rand.Seed(time.Now().UnixNano())

	// 简单的or-done-channel演示
	fmt.Println("=== 简单的or-done-channel示例 ===")
	done := make(chan struct{})

	// 生成一些数据
	intStream := generator(done, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)

	// 查找质数
	primeStream := primeFinder(done, intStream)

	// 打印前3个质数，然后取消操作
	fmt.Println("查找质数中...")
	for i := 0; i < 3; i++ {
		prime := <-primeStream
		fmt.Printf("找到质数: %d\n", prime)
	}

	fmt.Println("已找到足够的质数，取消操作")
	close(done)

	// 等待一会儿，让goroutine有时间响应取消信号
	time.Sleep(500 * time.Millisecond)

	// 使用超时机制的示例
	fmt.Println("=== 使用超时的or-done-channel示例 ===")

	// 搜索多个数据源，找到结果后取消其余搜索
	// 设置操作超时，防止资源浪费
	// 用户取消操作，立即停止后台工作
	// 优雅地关闭并发系统，确保所有资源都被释放

	// 创建一个3秒超时的上下文
	timoutDone, cancel := withTimeout(3 * time.Second)
	defer cancel()

	// 创建多个不同延时的worker
	w1 := worker(timoutDone, 1, 2*time.Second)
	w2 := worker(timoutDone, 2, 4*time.Second) // 这个会超时
	w3 := worker(timoutDone, 3, 5*time.Second) // 这个也会超时
	w4 := worker(timoutDone, 4, 1*time.Second)

	// 使用fanIn合并所有worker的输出
	mergedStream := fanIn(timoutDone, w1, w2, w3, w4)

	start := time.Now()

	// 收集所有未超时的结果
	var results []int
	for val := range mergedStream {
		if id, ok := val.(int); ok {
			results = append(results, id)
		}
	}

	fmt.Printf("收到的worker结果: %v (用时: %v)\n", results, time.Since(start))
	fmt.Println("注意：一些worker因超时而被取消")

	// 演示提前手动取消的情况
	fmt.Println("=== 手动取消的or-done-channel示例 ===")

	manualDone := make(chan struct{})

	// 创建多个worker
	workers := make([]<-chan interface{}, 5)
	for i := 0; i < 5; i++ {
		delay := time.Duration(rand.Intn(5)+1) * time.Second
		workers[i] = worker(manualDone, i+10, delay)
	}

	// 合并worker结果
	resultStream := fanIn(manualDone, workers...)

	// 只处理前2个结果，然后取消
	fmt.Println("仅等待前2个worker完成...")

	start = time.Now()

	// 收集2个结果后取消其余操作
	results = nil
	for i := 0; i < 2; i++ {
		val := <-resultStream
		if id, ok := val.(int); ok {
			results = append(results, id)
		}
	}

	fmt.Printf("收到2个结果: %v (用时: %v)\n", results, time.Since(start))
	fmt.Println("手动取消剩余worker...")

	close(manualDone)

	// 给goroutines时间退出
	time.Sleep(time.Second)
	fmt.Println("所有操作完成")
}
